// ============================================================================
// File: pkg/content/store.go - v2.0 with block-level deduplication
// ============================================================================
package content

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	"github.com/sirupsen/logrus"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/metrics"
)

type Store struct {
	root      string
	blockSize int64
	mu        sync.RWMutex
	blockMap  map[string][]string // hash -> file paths using this block
}

func NewStore(root string, blockSize int64) (*Store, error) {
	blobsDir := filepath.Join(root, "blobs", "sha256")
	blocksDir := filepath.Join(root, "blocks", "sha256")
	
	for _, dir := range []string{blobsDir, blocksDir} {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, err
		}
	}

	return &Store{
		root:      root,
		blockSize: blockSize,
		blockMap:  make(map[string][]string),
	}, nil
}

// Put stores data with file-level deduplication
func (s *Store) Put(data []byte) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	hash := sha256.Sum256(data)
	hashStr := hex.EncodeToString(hash[:])

	blobPath := s.blobPath(hashStr)
	if _, err := os.Stat(blobPath); err == nil {
		metrics.DedupSavedBytes.Add(float64(len(data)))
		metrics.DedupHitRate.WithLabelValues("file").Inc()
		return hashStr, nil
	}

	if err := os.MkdirAll(filepath.Dir(blobPath), 0755); err != nil {
		return "", err
	}

	if err := os.WriteFile(blobPath, data, 0444); err != nil {
		return "", err
	}

	metrics.StoredBytes.Add(float64(len(data)))
	metrics.DedupMissRate.WithLabelValues("file").Inc()
	return hashStr, nil
}

// PutWithBlockDedup stores large data with block-level deduplication
func (s *Store) PutWithBlockDedup(data []byte, path string) ([]string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if int64(len(data)) < s.blockSize {
		// Small file, use file-level dedup
		hash, err := s.putUnlocked(data)
		return []string{hash}, err
	}

	// Large file, use block-level dedup
	var blockHashes []string
	offset := 0
	totalSaved := int64(0)

	for offset < len(data) {
		end := offset + int(s.blockSize)
		if end > len(data) {
			end = len(data)
		}

		block := data[offset:end]
		hash := sha256.Sum256(block)
		hashStr := hex.EncodeToString(hash[:])

		blockPath := s.blockPath(hashStr)
		
		// Check if block exists (deduplication)
		if _, err := os.Stat(blockPath); err == nil {
			totalSaved += int64(len(block))
			metrics.DedupHitRate.WithLabelValues("block").Inc()
		} else {
			// New block, write it
			if err := os.MkdirAll(filepath.Dir(blockPath), 0755); err != nil {
				return nil, err
			}
			if err := os.WriteFile(blockPath, block, 0444); err != nil {
				return nil, err
			}
			metrics.StoredBytes.Add(float64(len(block)))
			metrics.DedupMissRate.WithLabelValues("block").Inc()
		}

		blockHashes = append(blockHashes, hashStr)
		s.blockMap[hashStr] = append(s.blockMap[hashStr], path)
		offset = end
	}

	if totalSaved > 0 {
		metrics.DedupSavedBytes.Add(float64(totalSaved))
		logrus.Debugf("block dedup saved %d bytes for %s", totalSaved, path)
	}

	return blockHashes, nil
}

func (s *Store) putUnlocked(data []byte) (string, error) {
	hash := sha256.Sum256(data)
	hashStr := hex.EncodeToString(hash[:])

	blobPath := s.blobPath(hashStr)
	if _, err := os.Stat(blobPath); err == nil {
		return hashStr, nil
	}

	if err := os.MkdirAll(filepath.Dir(blobPath), 0755); err != nil {
		return "", err
	}

	return hashStr, os.WriteFile(blobPath, data, 0444)
}

func (s *Store) Get(hash string) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// Try blob first
	data, err := os.ReadFile(s.blobPath(hash))
	if err == nil {
		return data, nil
	}

	// Try block
	return os.ReadFile(s.blockPath(hash))
}

func (s *Store) Exists(hash string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if _, err := os.Stat(s.blobPath(hash)); err == nil {
		return true
	}
	if _, err := os.Stat(s.blockPath(hash)); err == nil {
		return true
	}
	return false
}

func (s *Store) Link(hash, target string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	blobPath := s.blobPath(hash)
	if _, err := os.Stat(blobPath); err != nil {
		blobPath = s.blockPath(hash)
	}

	if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
		return err
	}

	return os.Link(blobPath, target)
}

func (s *Store) WriteFile(path string, data []byte, perm os.FileMode) error {
	blockHashes, err := s.PutWithBlockDedup(data, path)
	if err != nil {
		return err
	}

	// For files split into blocks, reconstruct them
	if len(blockHashes) > 1 {
		return s.reconstructFile(blockHashes, path, perm)
	}

	return s.Link(blockHashes[0], path)
}

func (s *Store) reconstructFile(blockHashes []string, target string, perm os.FileMode) error {
	f, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY, perm)
	if err != nil {
		return err
	}
	defer f.Close()

	for _, hash := range blockHashes {
		data, err := s.Get(hash)
		if err != nil {
			return fmt.Errorf("failed to get block %s: %w", hash, err)
		}
		if _, err := f.Write(data); err != nil {
			return err
		}
	}

	return nil
}

func (s *Store) blobPath(hash string) string {
	return filepath.Join(s.root, "blobs", "sha256", hash[:2], hash)
}

func (s *Store) blockPath(hash string) string {
	return filepath.Join(s.root, "blocks", "sha256", hash[:2], hash)
}

func (s *Store) GetStats() map[string]interface{} {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return map[string]interface{}{
		"block_size":        s.blockSize,
		"unique_blocks":     len(s.blockMap),
		"total_references":  s.countReferences(),
	}
}

func (s *Store) countReferences() int {
	total := 0
	for _, refs := range s.blockMap {
		total += len(refs)
	}
	return total
}

func HashReader(r io.Reader) (string, []byte, error) {
	h := sha256.New()
	data, err := io.ReadAll(r)
	if err != nil {
		return "", nil, err
	}

	h.Write(data)
	return hex.EncodeToString(h.Sum(nil)), data, nil
}
